#Run transformation.R first

#Working directory to file directory

setwd(dirname(rstudioapi::getActiveDocumentContext()$path))


#Create Plots folder, if it does not exist already
dir.create(file.path(getwd(), "Plots"))


#Install Required Packages


packages <- c("broom", "sjPlot", "estimatr", "tidyverse", "ggplot2", "patchwork", "stargazer", "svglite")

# Check if each package is installed, install if missing, and load it
for (pkg in packages) {
  if (!require(pkg, character.only = TRUE)) {
    install.packages(pkg)
    require(pkg, character.only = TRUE)
  }
}

#####
#Figure 1
#####

data = data %>% mutate(Level1F = as.factor(Level1),
                       Level2F = as.factor(Level2),
                       partisanshipF = as.factor(partisanship))

#Number of violations
M1 <- lm_robust(OutChoice ~ ViolNo, data = data, clusters = caseid)
M1part <- lm_robust(OutChoice ~ ViolNo + partisanship + ViolNo*partisanship, data = data, clusters = caseid)

#Types of violations
M2 <- lm_robust(OutChoice ~ Level1F + Level2F, data = data, clusters = caseid)
M2part <- lm_robust(OutChoice ~ Level1F + Level2F + partisanshipF + Level1F*partisanshipF, data = data, clusters = caseid)



#Models to Predictions
predictions <- predict(M1, newdata = data, se.fit = TRUE)
data$fit <- predictions$fit
data$se.fit <- predictions$se.fit
data$lower_ci <- data$fit - 1.96 * data$se.fit
data$upper_ci <- data$fit + 1.96 * data$se.fit

plot_data <- data %>%
  group_by(ViolNo) %>%
  summarise(
    Mean_Prediction = mean(fit),
    Lower_CI = mean(lower_ci),
    Upper_CI = mean(upper_ci),
    .groups = 'drop'
  )


pred_data0 <- expand.grid(
  Level1F = levels(data$Level1F),
  Level2F = "Compliance"
)

pred_data0$Predictions <- predict(M2, newdata = pred_data0, se.fit = TRUE)$fit
pred_data0$SE <- predict(M2, newdata = pred_data0, se.fit = TRUE)$se.fit
pred_data0$Upper <- pred_data0$Predictions + 1.96 * pred_data0$SE
pred_data0$Lower <- pred_data0$Predictions - 1.96 * pred_data0$SE


plot_data2 <- expand.grid(
  ViolNo = as.factor(c(0, 1, 2)), 
  partisanship = c("Democrat", "Republican")
)


predictions <- predict(M1part, newdata = plot_data2, se.fit = TRUE)

# Add the predictions and confidence intervals to plot_data2
plot_data2$predicted <- predictions$fit
plot_data2$conf.low <- predictions$fit - 1.96 * predictions$se.fit
plot_data2$conf.high <- predictions$fit + 1.96 * predictions$se.fit

# Now plot_data2 contains the predicted values and confidence intervals


pred_data <- expand.grid(
  partisanshipF = levels(data$partisanshipF),
  Level1F = levels(data$Level1F),
  Level2F = "Compliance"
)

pred_data$Predictions <- predict(M2part, newdata = pred_data, se.fit = TRUE)$fit
pred_data$SE <- predict(M2part, newdata = pred_data, se.fit = TRUE)$se.fit
pred_data$Upper <- pred_data$Predictions + 1.96 * pred_data$SE
pred_data$Lower <- pred_data$Predictions - 1.96 * pred_data$SE


#Individual plots
P1 <- ggplot(plot_data, aes(ViolNo, Mean_Prediction, fill = ViolNo)) + geom_bar(stat="identity", position=position_dodge(), alpha = 0.8) + geom_point(size = 2) +
  geom_errorbar(aes(ymin = Lower_CI, ymax = Upper_CI), width = .1, size = 1) + xlab("") + ylab("Defection Rates") +
  theme_minimal() +
  scale_x_discrete(breaks = c(0, 1, 2), labels = c("Compliance", "One Violation", "Two Violations")) +
  scale_fill_manual(values = c("0" = "lightblue3", "1" = "cornflowerblue", "2" = "navy")) + theme(legend.position = "none",
        text = element_text(family = "serif", size = 12)) + scale_y_continuous(breaks = c(0, .2, .4, .6), 
                                                                               labels = c("0 %", "20 %", "40 %", "60 %"), limits = c(0, .7)) +
  coord_flip()

P2 <- ggplot(pred_data0, aes(Level1F, Predictions)) + geom_bar(stat="identity", position=position_dodge(), alpha = 0.8, aes(fill = Level1F)) + geom_point(size = 2) +
  geom_errorbar(aes(ymin = Upper, ymax = Lower), width = .1, size = 1) + xlab("") + ylab("Defection Rates") +
  theme_minimal() +
  scale_x_discrete(breaks = c("Compliance", "Electoral", "FreeSpeech", "Judicial"), labels = c("Compliance", "Electoral Integrity", "Free Speech", "Judicial Integrity")) +
  scale_fill_manual(values = c("Electoral" = "navy", "FreeSpeech" = "cornflowerblue", "Judicial" = "navy", "Compliance" = "cornflowerblue")) + 
  theme(legend.position = "none",
        text = element_text(family = "serif", size = 12)) + scale_y_continuous(breaks = c(0, .2, .4, .6), 
                                                                                              labels = c("0 %", "20 %", "40 %", "60 %"), limits = c(0, .7)) +
    coord_flip()


P3 <- ggplot(plot_data2, aes(x = ViolNo, y = predicted, fill = partisanship)) +
  geom_bar(stat = "identity", position = position_dodge(width = 0.8), width = 0.7, alpha = 0.7) +
  geom_errorbar(aes(ymin = conf.low, ymax = conf.high), width = .1, size = 1, position = position_dodge(width = 0.8)) +
  geom_point(position = position_dodge(width = 0.8), size = 2) +
  labs(x = "", y = "Defection Rates", title = "") +
  theme_minimal() +  
  theme(legend.position = "null", text = element_text(family = "serif", size = 12)) +
  scale_x_discrete(breaks = c(0, 1, 2), labels = c("Compliance", "One Violation", "Two Violations")) +
  scale_fill_manual(values = c("Democrat" = "cornflowerblue", "Republican" = "lightblue3")) + 
  scale_y_continuous(breaks = c(0, .2, .4, .6), 
                     labels = c("0 %", "20 %", "40 %", "60 %")) +
  geom_text(data = subset(plot_data2, ViolNo != 0), 
            aes(label = partisanship, y = 0), 
            position = position_dodge(width = 0.8), 
            family = "serif", color = "black", size = 3.2, hjust = -0.5,
            fontface = "bold") + coord_flip()

P4 <- ggplot(pred_data, aes(x = Level1F, y = Predictions, fill = partisanshipF)) +
  geom_bar(stat = "identity", position = position_dodge(width = 0.8), width = 0.7, alpha = 0.7) +
  geom_errorbar(aes(ymin = Lower, ymax = Upper), width = .1, size = 1, position = position_dodge(width = 0.8)) +
  geom_point(position = position_dodge(width = 0.8), size = 2) +
  labs(x = "", y = "Defection Rates", title = "") +
  theme_minimal() +  
  theme(legend.position = "null", text = element_text(family = "serif", size = 12)) +
  scale_x_discrete(breaks = c("Compliance", "Electoral", "FreeSpeech", "Judicial"), 
                   labels = c("Compliance", "Electoral Integrity", "Free Speech", "Judicial Integrity")) +
  scale_fill_manual(values = c("Democrat" = "cornflowerblue", "Republican" = "lightblue3")) + 
  scale_y_continuous(breaks = c(0, .2, .4, .6), 
                     labels = c("0 %", "20 %", "40 %", "60 %")) +
    geom_text(data = subset(pred_data, Level1F != "Compliance"), 
            aes(label = partisanshipF, y = 0), 
            position = position_dodge(width = 0.8), 
            family = "serif", color = "black", size = 3.2, hjust = -0.5,
            fontface = "bold") + coord_flip()


#Figure 1

svg("plots/Figure1.svg", width = 8, height = 6)
P1 + P2 + P3 + P4 + plot_layout(heights = c(0.4, 0.6)) + plot_annotation(title = expression(italic("Biography and party labels only"))) &
  theme(plot.title = element_text(hjust = 0.5))
dev.off()

#####
#Figure 2
#####

#Government levels and violations
data = data %>% mutate(race = as.factor(race))
data$race <- factor(data$race, levels = c("Local", "State", "National"))
M3 <- lm_robust(OutChoice ~ race*ViolNo, data = data, clusters = caseid)

pred_data2 <- expand.grid(
  race = levels(data$race),
  ViolNo = levels(data$ViolNo)
)

pred_data2$Predictions <- predict(M3, newdata = pred_data2, se.fit = TRUE)$fit
pred_data2$SE <- predict(M3, newdata = pred_data2, se.fit = TRUE)$se.fit
pred_data2$Upper <- pred_data2$Predictions + 1.96 * pred_data2$SE
pred_data2$Lower <- pred_data2$Predictions - 1.96 * pred_data2$SE



P5 <- ggplot(pred_data2, aes(x = ViolNo, y = Predictions, fill = race)) +
  geom_bar(stat = "identity", position = position_dodge(width = 0.8), width = 0.7, alpha = 0.7) +
  geom_errorbar(aes(ymin = Lower, ymax = Upper), width = .1, size = 1, position = position_dodge(width = 0.8)) +
  geom_point(position = position_dodge(width = 0.8), size = 2) +
  labs(x = "", y = "Defection Rates", title = "") +
  theme_minimal() + theme(text = element_text(family = "serif", size = 12),
                          legend.position = "null") + 
  scale_x_discrete(breaks = c(0, 1, 2), labels = c("Compliance", "One Violation", "Two Violations")) +
  scale_fill_manual(values = c("Local" = "lightblue3", "State" = "cornflowerblue", "National" = "skyblue2")) + 
 scale_y_continuous(breaks = c(0, .2, .4, .6), labels = c("0 %", "20 %", "40 %", "60 %"), limits = c(0, .7)) +
  geom_text(data = subset(pred_data2, ViolNo != 0), 
            aes(label = race, y = 0), 
            position = position_dodge(width = 0.8), 
            family = "serif", color = "black", size = 3.2, hjust = -0.5,
            fontface = "bold") + coord_flip()

#Government Levels and Specific Violations
M4 <- lm_robust(OutChoice ~ Level1F + Level2F + race + Level1F*race, data = data, clusters = caseid)

pred_data3 <- expand.grid(
  race = levels(data$race),
  Level1F = levels(data$Level1F),
  Level2F = "Compliance"
)

pred_data3$Predictions <- predict(M4, newdata = pred_data3, se.fit = TRUE)$fit
pred_data3$SE <- predict(M4, newdata = pred_data3, se.fit = TRUE)$se.fit
pred_data3$Upper <- pred_data3$Predictions + 1.96 * pred_data3$SE
pred_data3$Lower <- pred_data3$Predictions - 1.96 * pred_data3$SE

P6 <- ggplot(pred_data3, aes(x = Level1F, y = Predictions, fill = race)) +
  geom_bar(stat = "identity", position = position_dodge(width = 0.8), width = 0.7, alpha = 0.7) +
  geom_errorbar(aes(ymin = Lower, ymax = Upper), width = .1, size = 1, position = position_dodge(width = 0.8)) +
  geom_point(position = position_dodge(width = 0.8), size = 2) +
  labs(x = "", y = "Defection Rates", title = "") +
  theme_minimal() +  theme(text = element_text(family = "serif", size = 12),
                           legend.position = "null") +
  scale_x_discrete(breaks = c("Compliance", "Electoral", "FreeSpeech", "Judicial"), labels = c("Compliance", "Electoral Integrity", "Free Speech", "Judicial Integrity")) +
  scale_fill_manual(values = c("Local" = "lightblue3", "State" = "cornflowerblue", "National" = "skyblue2")) + 
  scale_y_continuous(breaks = c(0, .2, .4, .6), 
                     labels = c("0 %", "20 %", "40 %", "60 %")) +
  geom_text(data = subset(pred_data3, Level1F != "Compliance"), 
            aes(label = race, y = 0), 
            position = position_dodge(width = 0.8), 
            family = "serif", color = "black", size = 3.2, hjust = -0.5,
            fontface = "bold") + coord_flip()


#Figure 2


svg("plots/Figure2.svg", height = 5, width = 8, family = "serif")
(P5 + P6) + plot_annotation(title = expression(italic("Biography and party labels only"))) &
  theme(plot.title = element_text(hjust = 0.5))
dev.off()

